"""
    基于楼盘信息写作，涉及楼盘户型信息（楼盘2）
"""
import datetime
import os
import random
import re
import time
import traceback
from typing import Dict, List

from langchain import PromptTemplate

try:
    from bot.openai_bot import OpenAIBot
    from conf.config import logger, BASE_DIR, config
    from db.building_dao import BuildingDAO
    from prompts.prompt_templates import BUILDING_INFO_BASE_PROMPT, BUILDING_GUIDE_PROMPT, ABSTRACT_AND_KEYWORDS_PROMPT, \
        TITLE_AND_ABSTRACT_AND_KEYWORDS_PROMPT, PARAGRAPH1_PROMPT, PARAGRAPH2_PROMPT, PARAGRAPH3_PROMPT
    from prompts.prompt_templates_within_building import BUILDING_GUIDE2_BASE, BUILDING_GUIDE2_PROMPT
    from utils.constants import EN_TO_CN, SALE_STATUS
    from utils.format_article import count_words, all_has_tags, remove_stop_words, remove_duplicate_paragraph_from_content, \
        extract_title_and_abstract_and_keywords
    from utils.laod_data import load_excel_sheet, load_building_id_list
    from utils.output_file import output_txt, output_json
except ModuleNotFoundError:
    import os
    import sys
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # 离开IDE也能正常导入自己定义的包
    from bot.openai_bot import OpenAIBot
    from conf.config import logger, BASE_DIR, config
    from db.building_dao import BuildingDAO
    from prompts.prompt_templates import BUILDING_INFO_BASE_PROMPT, BUILDING_GUIDE_PROMPT, ABSTRACT_AND_KEYWORDS_PROMPT, \
        TITLE_AND_ABSTRACT_AND_KEYWORDS_PROMPT, PARAGRAPH1_PROMPT, PARAGRAPH2_PROMPT, PARAGRAPH3_PROMPT
    from prompts.prompt_templates_within_building import BUILDING_GUIDE2_BASE, BUILDING_GUIDE2_PROMPT
    from utils.constants import EN_TO_CN, SALE_STATUS
    from utils.format_article import count_words, all_has_tags, remove_stop_words, remove_duplicate_paragraph_from_content, \
        extract_title_and_abstract_and_keywords
    from utils.laod_data import load_excel_sheet, load_building_id_list
    from utils.output_file import output_txt, output_json


def get_building_info(building_id: int) -> Dict:
    """
    获取楼盘信息
    :param building_id: 楼盘id
    :return: 楼盘信息
    """
    building_dao = BuildingDAO()
    building_info: Dict = building_dao.get_building_info(building_id=building_id)

    # TODO 2023-5-10 开盘时间一定小于交房时间，否则的话，运营答复需删除开盘时间。其他最好改数据库源头，此处先做临时处理
    start_date = building_info.get("BSTime")
    finish_date = building_info.get("BFTime")
    if start_date and finish_date:  # start_date和finish_date两者都不为None时，才能比较，否则还会引起转为datetime时还会报错
        start_date = datetime.datetime.strptime(str(start_date), '%Y-%m-%d').date()
        finish_date = datetime.datetime.strptime(str(finish_date), '%Y-%m-%d').date()
        if start_date >= finish_date:
            building_info.pop("BSTime")

    return building_info


def get_building_unit_type_list(building_id: int) -> List:
    """
    获取楼盘的户型列表
    :param building_id: 楼盘id
    :return:
    """
    building_dao = BuildingDAO()
    ls: List = building_dao.get_building_unit_type_list(building_id=building_id)

    ls_new = list()
    for index, item in enumerate(ls):
        unit_type_name = item.get("HName")  # 户型名
        room = item.get("Room")  # 室
        hall = item.get("Hall")  # 厅
        toilet = item.get("Toilet")  # 卫
        max_area = item.get("MaxArea")  # 建筑面积
        counter = item.get("counter")  # 首付几成

        if unit_type_name and room and max_area and (counter is not None):  # counter可以为0（零首付），但是不能为None
            item.update({"UnitType": f"{room}室{hall}厅{toilet}卫"})
            ls_new.append(item)

    return ls_new


def format_building_info(building_info: Dict, building_unit_type_list: List) -> str:
    """
    格式化楼盘信息
    an example::
        |   |   |
        | :--- | :--- |
        | 户型名 | a户型 |
        | 户型 | 3室2厅1卫 |
        | 建筑面积 | 约87平米 |
        | 参考总价 | 约186万/套 |
        | 参考首付 | 约55.8万 |

        ----

        |   |   |
        | :--- | :--- |
        | 户型名 | b户型 |
        | 户型 | 3室2厅1卫 |
        | 建筑面积 | 约87平米 |
        | 参考总价 | 约186万/套 |
        | 参考首付 | 约55.8万 |

    :param building_info: 楼盘信息
    :param building_unit_type_list: 楼盘户型列表
    :return: 格式化后的字符串

    """
    unit_type_info_string = ""
    average_price = building_info.get("AveragePrice")

    for index, item in enumerate(building_unit_type_list):
        build_area = item.get("MaxArea")  # 建筑面积
        counter = item.get("counter")  # 首付几成
        reference_total_payment = round(average_price * build_area / 10000)  # 参考总价
        reference_down_payment = round(reference_total_payment * counter / 10)  # 参考首付

        unit_type_info_string = ''.join([unit_type_info_string, BUILDING_GUIDE2_BASE,
                                         f'| {EN_TO_CN.get("HName")} | {item.get("HName")} |\n',
                                         f'| {EN_TO_CN.get("UnitType")} | {item.get("UnitType")} |\n',
                                         f'| {EN_TO_CN.get("MaxArea")} | 约{item.get("MaxArea")}平米 |\n',
                                         f'| {EN_TO_CN.get("ReferenceTotalPayment")} | 约{reference_total_payment}万/套 |\n',
                                         f'| {EN_TO_CN.get("ReferenceDownPayment")} | 约{reference_down_payment}万 |\n',
                                         '\n', '----', '\n\n'])

    return unit_type_info_string


def generate_article(prompt: str, filename: str, building_id: str, city_id: str, city_name: str):
    """
    生产文章
    :param prompt: 提交gpt的prompt
    :param filename: 保存的文件名
    :param building_id: 楼盘id
    :param city_id: 城市id
    :param city_name: 城市名
    :return:
    """
    data = dict()
    openai_bot = OpenAIBot(session_id=filename)
    prompt2 = PARAGRAPH1_PROMPT
    prompt3 = PARAGRAPH2_PROMPT
    prompt4 = PARAGRAPH3_PROMPT
    prompt5 = TITLE_AND_ABSTRACT_AND_KEYWORDS_PROMPT
    answer1 = openai_bot.reply(question=prompt, session_round=1)
    answer2 = openai_bot.reply(question=prompt2, session_round=2)
    answer3 = openai_bot.reply(question=prompt3, session_round=3)
    answer4 = openai_bot.reply(question=prompt4, session_round=4)
    answer5 = openai_bot.reply(question=prompt5, session_round=5)

    # 提取小标题
    pattern = r'<p><strong>.*?</strong></p>'
    subtitles = re.findall(pattern, answer1)
    # 正文段落
    paragraphs = [answer2, answer3, answer4]

    # 每段正文的第一段与小标题一致，则应该删除正文的第一段
    paragraphs = remove_duplicate_paragraph_from_content(subtitles=subtitles, paragraphs=paragraphs)

    # 带停用词的文章
    content_within_stop_words = "\n".join([subtitles[0], paragraphs[0], subtitles[1], paragraphs[1], subtitles[2], paragraphs[2], ])

    # 删除小标题的停用词
    for index, subtitle in enumerate(subtitles):
        subtitles[index] = remove_stop_words(text=subtitle, stop_words=config.get("stop_words_subtitle"))
    # 删除正文的停用词
    for index, paragraph in enumerate(paragraphs):
        paragraphs[index] = remove_stop_words(text=paragraph, stop_words=config.get("stop_words"))

    # 不带停用词的文章
    content = "\n".join([subtitles[0], paragraphs[0], subtitles[1], paragraphs[1], subtitles[2], paragraphs[2], ])

    title, abstract, keywords = extract_title_and_abstract_and_keywords(answer5)
    data.update({
        "id": str(filename), "title": title,
        "building_id": str(building_id), "city_id": str(city_id), "city_name": city_name,
        "abstract": abstract, "keywords": keywords, "content": content
    })

    # 判断摘要字数是否在限值内
    abstract_word_count = count_words(data.get("abstract"), excluded_char_ist=[])
    if abstract_word_count > config.get("word_limit_for_abstract"):
        raise ValueError(f"摘要字数超限值，共{abstract_word_count}字")

    # 判断正文字数是否足够
    num = count_words(data.get("content"), excluded_char_ist=config.get("excluded_char_ist"))
    if num <= config.get("word_limit_for_content").get("article_with_new_house_price"):
        raise ValueError(f"正文数字不足，只有{num}字")

    # 判断正文是否都有标签
    if not all_has_tags(text=data.get("content"), startswith_tags=config.get("startswith_tags"),
                        endswith_tags=config.get("endswith_tags")):
        raise ValueError(f"正文有些段落标签错误")

    # 输出到json(中文为不可读的Unicode编码)
    output_json(
        output_full_path=os.path.join(BASE_DIR, f"output/{filename}.json"),
        data=data, ensure_ascii=True
    )

    # prompt输出到txt文件
    output_txt(
        output_full_path=os.path.join(BASE_DIR, f'output/{filename}_prompt.txt'),
        txt="\n".join([prompt, prompt2, prompt3, prompt4, prompt5])
    )

    # 文章结果输出到txt文件
    text = "\n\n".join([answer1, answer2, answer3, answer4, answer5])
    count = count_words(text, excluded_char_ist=config.get("excluded_char_ist"))
    text = f"{text}\n\n==========\n正文字数统计：{count}"
    output_txt(
        output_full_path=os.path.join(BASE_DIR, f"output/{filename}.txt"),
        txt=text
    )

    # content_within_stop_words输出到txt文件
    output_txt(
        output_full_path=os.path.join(BASE_DIR, f'output/{filename}_content_within_stop_words.txt'),
        txt=content_within_stop_words
    )


def main():
    fail_list = list()
    building_id_list = load_building_id_list()
    for item in building_id_list:
        building_id = item.get("building_id")
        building_name = item.get("building_name")
        filename = item.get("filename")
        logger.info(f"文件名为{filename}，楼盘id为{building_id}的文章正在进行创作...")
        try:
            # 查询数据库获取楼盘信息
            building_info = get_building_info(building_id=building_id)
            city_id = building_info.get("city_id")
            city_name = building_info.get("city_name")
            # 查询楼盘的户型信息
            building_unit_type_list = get_building_unit_type_list(building_id=building_id)
            if not building_unit_type_list:
                logger.warning(f"查询到的户型信息为0，跳过生产！")
                continue

            # 将楼盘信息格式化为字符串
            building_info_string = format_building_info(building_info, building_unit_type_list)

            # 渲染模板
            prompt = PromptTemplate(
                input_variables=["building_info", "date", "building_name"],
                template=BUILDING_GUIDE2_PROMPT,
                template_format="jinja2"
            )
            prompt_formatted = prompt.format(building_info=building_info_string,
                                             date=datetime.date.today().strftime('%Y-%m-%d'),
                                             building_name=building_name)

            generate_article(prompt=prompt_formatted, filename=f'{filename}',
                             building_id=building_id, city_id=city_id, city_name=city_name)
            time.sleep(random.uniform(10, 15))

        except Exception:
            error_str = traceback.format_exc()
            logger.error(error_str)
            fail_list.append({"filename": filename, "log": error_str})

    logger.warning(f'失败列表：\n{[item.get("filename") for item in fail_list]}')
    logger.warning(f'失败列表详情：\n{fail_list}')


if __name__ == '__main__':
    main()
